import matplotlib.pyplot as plt
import numpy as np

def tensor_debug_view(tensor, title="Tensor Preview", denormalize=True, save_path=None):
    """
    可视化任意图像张量
    
    参数:
        tensor: 输入张量 (PyTorch/TensorFlow/numpy)
        title: 图像标题 (默认: "Tensor Preview")
        denormalize: 是否自动反归一化 (默认: True)
        save_path: 保存图像路径 (默认: 仅显示)
    """
    # 转换到 numpy 数组
    if hasattr(tensor, 'numpy'):
        arr = tensor.numpy()
    elif hasattr(tensor, 'cpu'):
        arr = tensor.cpu().numpy()
    else:
        arr = np.array(tensor)
    
    # 处理不同维度格式
    if arr.ndim == 4:  # [B, C, H, W] 或 [B, H, W, C]
        arr = arr[0]  # 取批次中第一张图
    
    if arr.ndim == 3:
        # 检测通道位置 (PyTorch: CHW vs TensorFlow: HWC)
        if arr.shape[0] < arr.shape[-1] and arr.shape[0] <= 4:
            arr = np.transpose(arr, (1, 2, 0))  # CHW -> HWC
    
    # 自动反归一化处理
    if denormalize:
        if arr.min() < 0:  # 假设范围 [-1, 1]
            arr = (arr + 1) / 2.0
        elif arr.max() <= 1 and arr.min() >= 0:  # 范围 [0, 1]
            arr = arr * 255
    
    # 转换到整数类型
    if arr.dtype != np.uint8:
        arr = np.clip(arr, 0, 255).astype(np.uint8)
    
    # 处理单通道图像
    if arr.ndim == 2:
        cmap = 'gray'
    else:
        cmap = None
    
    # 可视化
    plt.figure(figsize=(8, 6))
    plt.imshow(arr, cmap=cmap)
    plt.title(f"{title}\nShape: {tensor.shape} | Type: {type(tensor)}")
    plt.axis('off')
    
    # 添加像素值范围信息
    plt.figtext(0.5, 0.01, 
                f"Value Range: [{tensor.min().item():.3f}, {tensor.max().item():.3f}]",
                ha="center")
    
    if save_path:
        plt.savefig(save_path, bbox_inches='tight')
        print(f"➤ 图像已保存至: {save_path}")
    else:
        plt.show()